import cv2
import apriltag
import os
import time
import json
import sys
sys.path.append("../")

from typing import Tuple
from arm_control.angle import ArmAngle
from arm_control.move import ArmMove

"""
读取保存为 JSON 格式的配置文件，并返回一个字典。
输入:
    file (str): 配置文件的路径。
输出：
    dict: 包含配置信息的字典。
"""
def load_config(file):
     # 检查文件是否存在
    if not os.path.isfile(file):
        raise FileNotFoundError(f"配置文件 {file} 不存在.")
     # 读取文件内容
    with open(file, 'r', encoding='utf-8') as f:
        config = json.load(f)
    return config


################################## 载入相关文件 #############################
CAM_CONFIG_FILE = os.path.expanduser("~/.bc_config/cam_calibration.config")
# 读取配置文件
cam_cfg = load_config(CAM_CONFIG_FILE)
# 将配置文件内容载入变量
# 机械臂基座的 x 坐标
arm_base_coordinate_x = cam_cfg["arm_base_coordinate"]["x"]
# 机械臂基座的 y 坐标
arm_base_coordinate_y = cam_cfg["arm_base_coordinate"]["y"]
# 像素密度
pixel_per_cm_ratio = cam_cfg["pixel_per_cm_ratio"]

# 是否执行抓取
arm_action = False

"""
计算像素坐标系下的点距离原点的绝对距离
注意：这里的坐标系是以摄像机为原点,向右为x正方向,向下为y正方向方向.
输入:
    pixel_x (int): 像素坐标系下的x坐标
    pixel_y (int): 像素坐标系下的y坐标
    pixel_per_cm_ratio (float): 像素与厘米的比值, 1cm = ?pixel
输出:
tuple (float, float): x方向和y方向上的距离(单位: cm)
"""
def compute_pixel_to_dist(
    pixel_x: int, pixel_y: int, pixel_per_cm_ratio: float
) -> Tuple[float, float]:
    dist_x, dist_y = pixel_x / pixel_per_cm_ratio, pixel_y / pixel_per_cm_ratio
    return round(dist_x, 2), round(dist_y, 2)



if __name__=="__main__":
    cap = cv2.VideoCapture(0)
    cv2.namedWindow("Video Frame")

    # 定义target_position
    target_position = None

    # 创建 ArmMove and ArmAngle
    arm_move = ArmMove()
    arm_angle = ArmAngle(base_pos=[0,0,0])
    
    ############### Part1 初始化 Aprialtag ##############
    # 创建Apriltag检测器
    detector = apriltag.Detector()

    #####################################################################
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        #####################Part2 填入检测Apriltag的代码 #####################
        # 填写到Part2  检测Apriltag的代码
        gray_image = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        tags = detector.detect(gray_image)

        # 我们假定视角中只有一个Apriltag
        for tag in tags:
            # 绘制标签边界
            for idx in range(len(tag.corners)):
                pt1 = tuple(tag.corners[idx - 1].astype(int))
                pt2 = tuple(tag.corners[idx].astype(int))
                cv2.line(frame, pt1, pt2, (0, 255, 0), 2)

            # 打印标签的ID和中心点位置
            #print(f'ID: {tag.tag_id}, Center: {tag.center}')

            # 保存中心点位置为 target_position
            target_position = tag.center
            
        #################################################################
        # 当用户按下`s`， 执行机械臂抓取， 否则不进行抓取。
        if cv2.waitKey(1) & 0xFF == ord('s'):
            arm_action = True

        #####################Part3  填入机械臂控制的代码  #####################
        if arm_action:
            arm_action = False

            # 返回的Apriltag的格式为[float, float]， 我们需要将其转换为[int, int]，图像的像素个数应该是整数。
            target_position = [int(x) for x in target_position]

            # 基座中心坐标
            center_coordinates = (arm_base_coordinate_x, arm_base_coordinate_y)
            
            # 基座中心到待抓取的AprilTag中心的像素距离
            diff_x_pixel = abs(target_position[0] - center_coordinates[0])
            diff_y_pixel = abs(target_position[1] - center_coordinates[1])
            
            # 像素距离转化到真实物理距离
            distance_x, distance_y = compute_pixel_to_dist(diff_x_pixel, diff_y_pixel, pixel_per_cm_ratio)
            
            print(f"x: {distance_x}, y: {distance_y}")

            # 对于在高度上的距离, 基座中心点的高度为0,  Apriltag的高度为立方体的边长，所以抓取时吸盘高度应该略小于Apriltag高度。约等于 3.9
            distance_z = 3.9

            # 明确起始点
            grab_position = [distance_x, distance_y, distance_z]
            grab_position_h = [distance_x, distance_y, distance_z+4.0]

            # 为方便起见， 定义放置点与起始点沿机械臂基座中心点yz平面对称
            release_position = [-distance_x, distance_y, distance_z]
            release_position_h = [-distance_x, distance_y, distance_z+4.0]

            # 获取机械臂位于起始点上方时, 对应的关节电机位置
            steps = arm_angle.get_steps(*grab_position_h)
            # 机械臂电机执行到对应位置
            arm_move.go_to(*steps)
            # 获取机械臂位于起始点时, 对应的关节电机位置
            steps = arm_angle.get_steps(*grab_position)
            # 机械臂电机执行到对应位置
            arm_move.go_to(*steps)
            # 机械臂抓取
            arm_move.grip()

            time.sleep(0.1)
            
            # 获取机械臂位于放置点上方时, 对应的关节电机位置
            steps = arm_angle.get_steps(*release_position_h)
            # 机械臂电机执行到对应位置
            arm_move.go_to(*steps)
            # 获取机械臂位于放置时, 对应的关节电机位置
            steps = arm_angle.get_steps(*release_position)
            # 机械臂电机执行到对应位置
            arm_move.go_to(*steps)
            # 机械臂放置
            arm_move.loose()
            
            time.sleep(1)
            
            arm_move.reset()

        cv2.imshow("Video Frame", frame)
        
        if cv2.waitKey(10) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()
